package org.example;

import ai.djl.Device;
import ai.djl.engine.Engine;
import ai.djl.util.cuda.CudaUtils;
import java.lang.management.MemoryUsage;
import java.util.ArrayList;
import java.util.List;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class Main {
    public static void main(String[] args) {
        System.out.println("GPU count: " + Engine.getInstance().getGpuCount());

        var devices = tryAllGpus();
        devices.forEach(it ->
                System.out.println("Device : " + it.getDeviceType() + " " + it.getDeviceId())
        );

        // 查询显存
        var gpuDevice = devices.stream().filter(Device::isGpu).findFirst().get();
        MemoryUsage mem = CudaUtils.getGpuMemory(gpuDevice);
        System.out.println("显存:" + mem.getMax());
    }


    public static List<Device> tryAllGpus() {
        var devices = new ArrayList<Device>();
        int gpuCount = Engine.getInstance().getGpuCount();
        if (gpuCount > 0) {
            for (int i = 0; i < gpuCount; i++) {
                devices.add(Device.gpu(i));
            }
        }

        if (devices.isEmpty()) {
            devices.addAll(List.of(Engine.getInstance().getDevices()));
        }

        if (!devices.stream().anyMatch(Device::isGpu)) {
            try {
                devices.add(Device.gpu());
            } catch (Exception e) {
                log.warn("***找不到显卡设备***", e);
            }
        }

        return devices;
    }
}
